import pathlib

import matplotlib.pyplot
import torch

from checkpoint import Checkpoint
from checkpoint_manager import CheckpointManager
from scaler import Scaler
from lotka_volterra_dataset import LotkaVolterraDataset
from lvminn import Lvminn


def main():
    dataset = LotkaVolterraDataset.load(pathlib.Path("./checkpoints/data.pt"))
    t_min = float(dataset.t_min())
    t_max = float(dataset.t_max())

    _, checkpoint = (CheckpointManager(pathlib.Path("./checkpoints"))
                     .last_existed_loaded(Checkpoint.load))
    assert checkpoint is not None
    model = Lvminn(Scaler(t_min, t_max))
    model.load_state_dict(checkpoint.model_state())
    model.eval()

    matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    matplotlib.pyplot.scatter(dataset.t1(), dataset.x1(), label="Data $x_1$")
    matplotlib.pyplot.scatter(dataset.t2(), dataset.x2(), label="Data $x_2$")

    t = torch.linspace(t_min, t_max, 1000)
    x = torch.detach(model(t))
    matplotlib.pyplot.plot(t, x[:, 0], label="Lvminn $x_1$")
    matplotlib.pyplot.plot(t, x[:, 1], label="Lvminn $x_2$")

    matplotlib.pyplot.legend()
    matplotlib.pyplot.xlabel(r"$t$")
    matplotlib.pyplot.ylabel(r"$x$")

    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
